from fastapi import FastAPI, UploadFile, File, Form
from fastapi.middleware.cors import CORSMiddleware
from rag_chain import process_document
import uvicorn
import requests
from datetime import datetime
import json
import os
from dotenv import load_dotenv
import time
from requests.adapters import HTTPAdapter
from urllib3.util import Retry
from rag_chain import RagChatBot
import logging
from langchain.text_splitter import RecursiveCharacterTextSplitter
import re

# 配置日志记录
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s',
                   handlers=[
                       logging.FileHandler(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'main.log')),
                       logging.StreamHandler()
                   ])
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

app = FastAPI()

# 在这里定义一个变量来存储当前激活的向量存储
# 注意：这只是一个简单的实现，实际应用可能需要更复杂的会话或用户管理来处理不同的向量存储
current_vectorstore = None

# 天气API配置
OPENWEATHER_API_KEY = os.getenv("OPENWEATHER_API_KEY")
OPENWEATHER_API_URL = "http://api.openweathermap.org/data/2.5/weather"
OPENWEATHER_FIND_URL = "http://api.openweathermap.org/geo/1.0/direct"
OPENWEATHER_FORECAST_URL = "http://api.openweathermap.org/data/2.5/forecast"

# 城市名称映射（中文到英文）
CITY_NAME_MAPPING = {
    # 直辖市
    "北京": "Beijing",
    "上海": "Shanghai",
    "天津": "Tianjin",
    "重庆": "Chongqing",
    
    # 省会城市
    "广州": "Guangzhou",
    "杭州": "Hangzhou",
    "南京": "Nanjing",
    "武汉": "Wuhan",
    "成都": "Chengdu",
    "西安": "Xian",
    "长沙": "Changsha",
    "郑州": "Zhengzhou",
    "济南": "Jinan",
    "福州": "Fuzhou",
    "南昌": "Nanchang",
    "合肥": "Hefei",
    "石家庄": "Shijiazhuang",
    "太原": "Taiyuan",
    "呼和浩特": "Hohhot",
    "沈阳": "Shenyang",
    "长春": "Changchun",
    "哈尔滨": "Harbin",
    "海口": "Haikou",
    "南宁": "Nanning",
    "贵阳": "Guiyang",
    "昆明": "Kunming",
    "拉萨": "Lhasa",
    "兰州": "Lanzhou",
    "西宁": "Xining",
    "银川": "Yinchuan",
    "乌鲁木齐": "Urumqi",
    
    # 计划单列市
    "深圳": "Shenzhen",
    "大连": "Dalian",
    "青岛": "Qingdao",
    "宁波": "Ningbo",
    "厦门": "Xiamen",
    
    # 其他主要城市
    "苏州": "Suzhou",
    "无锡": "Wuxi",
    "常州": "Changzhou",
    "南通": "Nantong",
    "扬州": "Yangzhou",
    "镇江": "Zhenjiang",
    "泰州": "Taizhou",
    "徐州": "Xuzhou",
    "连云港": "Lianyungang",
    "淮安": "Huaian",
    "盐城": "Yancheng",
    "宿迁": "Suqian",
    "温州": "Wenzhou",
    "嘉兴": "Jiaxing",
    "湖州": "Huzhou",
    "绍兴": "Shaoxing",
    "金华": "Jinhua",
    "衢州": "Quzhou",
    "舟山": "Zhoushan",
    "台州": "Taizhou",
    "丽水": "Lishui",
    "九江": "Jiujiang",
    "赣州": "Ganzhou",
    "吉安": "Jian",
    "宜春": "Yichun",
    "抚州": "Fuzhou",
    "上饶": "Shangrao",
    "景德镇": "Jingdezhen",
    "萍乡": "Pingxiang",
    "新余": "Xinyu",
    "鹰潭": "Yingtan",
    "襄阳": "Xiangyang",
    "宜昌": "Yichang",
    "荆州": "Jingzhou",
    "荆门": "Jingmen",
    "鄂州": "Ezhou",
    "孝感": "Xiaogan",
    "黄冈": "Huanggang",
    "咸宁": "Xianning",
    "随州": "Suizhou",
    "恩施": "Enshi",
    "仙桃": "Xiantao",
    "潜江": "Qianjiang",
    "天门": "Tianmen",
    "神农架": "Shennongjia"
}

# 初始化 RagChatBot
rag_bot = RagChatBot()

def get_english_city_name(chinese_city):
    """将中文城市名转换为英文城市名"""
    return CITY_NAME_MAPPING.get(chinese_city, chinese_city)

def get_supported_cities():
    """获取所有支持的城市列表"""
    return list(CITY_NAME_MAPPING.keys())

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 创建带有重试机制的会话
def create_session():
    session = requests.Session()
    retry_strategy = Retry(
        total=3,  # 最大重试次数
        backoff_factor=1,  # 重试间隔
        status_forcelist=[500, 502, 503, 504]  # 需要重试的HTTP状态码
    )
    adapter = HTTPAdapter(max_retries=retry_strategy)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session

def find_city(city_name):
    """使用OpenWeather API查找城市"""
    try:
        session = create_session()
        params = {
            'q': city_name,
            'limit': 1,
            'appid': OPENWEATHER_API_KEY
        }
        response = session.get(OPENWEATHER_FIND_URL, params=params, timeout=10)
        response.raise_for_status()
        data = response.json()
        
        if data and len(data) > 0:
            return data[0]
        return None
    except requests.exceptions.Timeout:
        return None
    except requests.exceptions.RequestException as e:
        print(f"查找城市时发生错误: {str(e)}")
        return None

def get_weather(city, days=1):
    """获取天气信息，支持多天预报（最多5天）"""
    if not OPENWEATHER_API_KEY:
        return "未配置天气API密钥"
    try:
        city_data = find_city(city)
        if not city_data:
            return f"抱歉，找不到城市 '{city}' 的天气信息，请确认城市名称是否正确。"
        session = create_session()
        if days == 1:
            # 单天实时天气
        params = {
            'lat': city_data['lat'],
            'lon': city_data['lon'],
            'appid': OPENWEATHER_API_KEY,
                'units': 'metric',
                'lang': 'zh_cn'
        }
        response = session.get(OPENWEATHER_API_URL, params=params, timeout=10)
        response.raise_for_status()
        data = response.json()
        weather_info = {
                'city': city,
            'temperature': f"{data['main']['temp']}°C",
            'weather': data['weather'][0]['description'],
            'humidity': f"{data['main']['humidity']}%",
            'wind_speed': f"{data['wind']['speed']} m/s",
            'update_time': datetime.fromtimestamp(data['dt']).strftime('%Y-%m-%d %H:%M:%S')
        }
        weather_text = f"""
城市：{weather_info['city']}
温度：{weather_info['temperature']}
天气：{weather_info['weather']}
湿度：{weather_info['humidity']}
风速：{weather_info['wind_speed']}
更新时间：{weather_info['update_time']}
"""
        return weather_text
        else:
            # 多天预报
            params = {
                'lat': city_data['lat'],
                'lon': city_data['lon'],
                'appid': OPENWEATHER_API_KEY,
                'units': 'metric',
                'lang': 'zh_cn'
            }
            response = session.get(OPENWEATHER_FORECAST_URL, params=params, timeout=10)
            response.raise_for_status()
            data = response.json()
            from collections import defaultdict
            daily = defaultdict(list)
            for item in data['list']:
                date = item['dt_txt'][:10]
                hour = int(item['dt_txt'][11:13])
                if 11 <= hour <= 13:
                    daily[date].append(item)
            days_list = sorted(daily.keys())[:days]
            result = f"城市：{city}\n"
            for d in days_list:
                if daily[d]:
                    item = daily[d][0]
                    temp = item['main']['temp']
                    desc = item['weather'][0]['description']
                    result += f"{d} 天气：{desc}，温度：{temp}°C\n"
            return result
    except Exception as e:
        return f"获取天气信息失败: {str(e)}"

def process_web_search(query):
    """处理联网搜索请求"""
    try:
        # 检查是否是时间或日期查询
        time_keywords = ["现在时间", "当前时间", "几点了", "what time", "current time"]
        date_keywords = ["今天几号", "今天是几号", "日期", "今天什么日期", "今天多少号"]
        if any(keyword in query.lower() for keyword in time_keywords):
            from datetime import datetime
            import pytz
            beijing_tz = pytz.timezone('Asia/Shanghai')
            current_time = datetime.now(beijing_tz)
            return f"当前北京时间是：{current_time.strftime('%Y-%m-%d %H:%M:%S')}"
        if any(keyword in query for keyword in date_keywords):
            from datetime import datetime
            import pytz
            beijing_tz = pytz.timezone('Asia/Shanghai')
            current_time = datetime.now(beijing_tz)
            return f"今天是：{current_time.strftime('%Y-%m-%d')}"
        # 检查是否是天气查询
        if "天气" in query or "怎么样" in query or "如何" in query or "怎样" in query:
            # 检查天数
            import re
            days = 1
            if "明天" in query:
                days = 2
            elif "后天" in query:
                days = 3
            else:
                match = re.search(r"(\d+)天", query)
                if match:
                    days = int(match.group(1))
                    if days < 1:
                        days = 1
                    if days > 5:
                        days = 5
            # 联网搜索时，优先用城市名映射表做模糊匹配
            city = None
            for cname in CITY_NAME_MAPPING.keys():
                if cname in query:
                    city = cname
                    break
            if not city:
                city = query
            return get_weather(city, days=days)
        # 其他类型的联网搜索
        return "暂不支持该类型的联网搜索"
    except Exception as e:
        return f"处理联网搜索请求时发生错误：{str(e)}"

@app.post("/upload/")
async def upload_document(file: UploadFile = File(...)):
    try:
        # 读取文件内容
        file_content = await file.read()
        file_text = file_content.decode('utf-8')

        # 处理文档并创建向量存储
        try:
            rag_bot.process_document(file_text)
            # 验证向量存储是否成功创建
            if not rag_bot.vectorstore:
                return {"status": "error", "message": "Failed to create vectorstore"}
            return {
                "status": "success",
                "message": f"File '{file.filename}' processed and indexed successfully.",
                "filename": file.filename
            }
        except Exception as e:
            logger.error(f"Error processing document: {str(e)}")
            return {"status": "error", "message": f"Error processing file: {str(e)}"}
    except Exception as e:
        logger.error(f"Error in upload endpoint: {str(e)}")
        return {"status": "error", "message": f"Error uploading file: {str(e)}"}

@app.post("/chat")
async def chat_api(
    query: str = Form(...), 
    deep_search: bool = Form(False), 
    reasoning: bool = Form(False), 
    web_search: bool = Form(False),
    rag_search: bool = Form(False),
    thread_id: str = Form("default_thread")
):
    try:
        if web_search:
            response_content = process_web_search(query)
        elif rag_search:
            # 检查向量存储是否存在并且可用
            if not rag_bot.vectorstore:
                logger.warning("RAG search requested but no vectorstore available")
                return {"response": "请先上传知识库文件才能使用RAG检索功能。"}
            
            # 使用 RAG 检索模式
            try:
                response_content = rag_bot.chat(
                    query=query,
                    deep_search=deep_search,
                    reasoning=reasoning
                )
                if not response_content or response_content.strip() == "":
                    response_content = "在知识库中未找到相关信息。"
            except Exception as e:
                logger.error(f"Error in RAG chat: {str(e)}")
                return {"response": "RAG检索过程中发生错误，请重试。"}
        else:
            # 普通聊天模式
            response_content = rag_bot.chat(
                query=query,
                deep_search=deep_search,
                reasoning=reasoning
            )
        return {"response": response_content}
    except Exception as e:
        logger.error(f"Error in chat endpoint: {str(e)}")
        return {"response": f"处理聊天请求时发生错误: {str(e)}"}

def extract_city(query):
    # 匹配第一个连续的中文字符串
    match = re.search(r'[\u4e00-\u9fa5]+', query)
    if match:
        return match.group(0)
    return query.strip()

if __name__ == "__main__":
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)
    # 后端使用 FastAPI 接收请求，调用相应的处理逻辑（如 RAG 流程和文档处理），并将结果通过 HTTP 响应返回给前端。
    # uvicorn main:app --reload --port 8000